package com.bruce.ai.alibaba.workflow.controller;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import com.alibaba.cloud.ai.graph.CompileConfig;
import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;
import com.alibaba.cloud.ai.graph.store.StoreItem;
import com.alibaba.cloud.ai.graph.store.stores.MemoryStore;

import jakarta.annotation.Resource;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * 长期记忆工作流
 */
@RestController
@RequestMapping("/long-memory")
public class LongMemoryWorkflowController {


    private CompiledGraph memoryWorkflowCompiledGraph ;

    @Resource
    private MemoryStore memoryStore;



    private static final Logger logger = LoggerFactory.getLogger(ShortMemoryWorkflowController.class);


    public LongMemoryWorkflowController(@Qualifier("longMemoryWorkflow") StateGraph longMemoryWorkflow) throws Exception {
        this.memoryWorkflowCompiledGraph = longMemoryWorkflow
                .compile(CompileConfig.builder().build());
    }

    /*
     * 展开工作流
     */
    @RequestMapping(value ="/expand", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> expand(@RequestParam(name = "userId",defaultValue = "user_0001") String userId) throws Exception{
        return Mono.defer(() -> Mono.justOrEmpty(memoryWorkflowCompiledGraph.invoke(Map.of("userId", userId), getRunnableConfig(userId)))
                .map(OverAllState::data)
                .map(data -> ServerSentEvent.builder(data.toString()).build()))
            .flux()
            .doOnCancel(() -> logger.info("Client disconnected from stream"))
            .doOnError(e -> logger.error("Error occurred during streaming", e));
    }


    private RunnableConfig getRunnableConfig(String userId) {
        // 定义检查点配置 内存检查点
        Map<String,Object> profileData = new HashMap<>() ;
        profileData.put("name", "bruce");
        profileData.put("preferences", "喜欢唱跳、rap、篮球");
        StoreItem profileItem = StoreItem.of(List.of("user_profiles"),"user_0001", profileData);
        memoryStore.putItem(profileItem);
        return RunnableConfig.builder()
                .threadId(userId)
                .store(memoryStore)
                .build();
    }

}
